// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <nebula/future/serial_executor.h>
#include <nebula/future/mutex.h>
#include <mutex>
#include <string>
#include <thread>
#include <condition_variable>

namespace nebula::internal {

    struct SerialExecutor::State {
        std::deque<Task> task_queue;
        std::mutex mutex;
        std::condition_variable wait_for_tasks;
        std::thread::id current_thread;
        bool paused{false};
        bool finished{false};
    };

    SerialExecutor::SerialExecutor() : state_(std::make_shared<State>()) {
    }

    SerialExecutor::~SerialExecutor() {
        auto state = state_;
        std::unique_lock<std::mutex> lk(state->mutex);
        if (!state->task_queue.empty()) {
            // We may have remaining tasks if the executor is being abandoned.  We could have
            // resource leakage in this case.  However, we can force the cleanup to happen now
            state->paused = false;
            lk.unlock();
            RunLoop();
            lk.lock();
        }
    }

    int SerialExecutor::GetNumTasks() {
        auto state = state_;
        return static_cast<int>(state_->task_queue.size());
    }

    turbo::Status SerialExecutor::spawn_real(TaskHints hints, turbo::FnOnce<void()> task,
                                            StopToken stop_token, StopCallback &&stop_callback) {
#ifdef NEBULA_WITH_OPENTELEMETRY
        // Wrap the task to propagate a parent tracing span to it
            // XXX should there be a generic utility in tracing_internal.h for this?
            task = [func = std::move(task),
                    active_span =
                        ::nebula::internal::tracing::GetTracer()->GetCurrentSpan()]() mutable {
              auto scope = ::nebula::internal::tracing::GetTracer()->WithActiveSpan(active_span);
              std::move(func)();
            };
#endif
        // While the SerialExecutor runs tasks synchronously on its main thread,
        // spawn_real may be called from external threads (e.g. when transferring back
        // from blocking I/O threads), so we need to keep the state alive *and* to
        // lock its contents.
        //
        // Note that holding the lock while notifying the condition variable may
        // not be sufficient, as some exit paths in the main thread are unlocked.
        auto state = state_;
        {
            std::lock_guard<std::mutex> lk(state->mutex);
            if (state_->finished) {
                return turbo::invalid_argument_error(
                        "Attempt to schedule a task on a serial executor that has already finished or "
                        "been abandoned");
            }
            state->task_queue.push_back(
                    Task{std::move(task), std::move(stop_token), std::move(stop_callback)});
        }
        state->wait_for_tasks.notify_one();
        return turbo::OkStatus();
    }

    void SerialExecutor::finish() {
        auto state = state_;
        {
            std::lock_guard<std::mutex> lk(state->mutex);
            state->finished = true;
        }
        state->wait_for_tasks.notify_one();
    }


    void SerialExecutor::Pause() {
        // Same comment as spawn_real above
        auto state = state_;
        {
            std::lock_guard<std::mutex> lk(state->mutex);
            state->paused = true;
        }
        state->wait_for_tasks.notify_one();
    }

    bool SerialExecutor::IsFinished() {
        std::lock_guard<std::mutex> lk(state_->mutex);
        return state_->finished;
    }

    void SerialExecutor::Unpause() {
        auto state = state_;
        {
            std::lock_guard<std::mutex> lk(state->mutex);
            state->paused = false;
        }
    }

    bool SerialExecutor::owns_this_thread() {
        std::lock_guard lk(state_->mutex);
        return std::this_thread::get_id() == state_->current_thread;
    }


    void SerialExecutor::RunLoop() {
        // This is called from the SerialExecutor's main thread, so the
        // state is guaranteed to be kept alive.
        std::unique_lock<std::mutex> lk(state_->mutex);
        state_->current_thread = std::this_thread::get_id();
        // If paused we break out immediately.  If finished we only break out
        // when all work is done.
        while (!state_->paused && !(state_->finished && state_->task_queue.empty())) {
            // The inner loop is to check if we need to sleep (e.g. while waiting on some
            // async task to finish from another thread pool).  We still need to check paused
            // because sometimes we will pause even with work leftover when processing
            // an async generator
            while (!state_->paused && !state_->task_queue.empty()) {
                Task task = std::move(state_->task_queue.front());
                state_->task_queue.pop_front();
                lk.unlock();
                if (!task.stop_token.is_stop_requested()) {
                    std::move(task.callable)();
                } else {
                    if (task.stop_callback) {
                        std::move(task.stop_callback)(task.stop_token.Poll());
                    }
                    // Can't break here because there may be cleanup tasks down the chain we still
                    // need to run.
                }
                lk.lock();
            }
            // In this case we must be waiting on work from external (e.g. I/O) executors.  Wait
            // for tasks to arrive (typically via transferred futures).
            state_->wait_for_tasks.wait(lk, [&] {
                return state_->paused || state_->finished || !state_->task_queue.empty();
            });
        }
        state_->current_thread = {};
    }

}  // namespace nebula::internal
